'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: agent_queue_manager.py
* @Time: 2025/10/9
* @All Rights Reserve By Brtc
'''
import queue
import time
import uuid
from enum import Enum
from queue import Queue
from typing import Generator
from uuid import UUID

from redis import Redis

from internal.entity.conversation_entity import InvokeFrom
from internal.core.agent.entities.queue_entity import AgentThought, QueueEvent


class PublishFrom(str, Enum):
    """消息队列发布来源, 用于记录队列管理器中的发布来源"""
    APPLICATION_MANAGER = 1,
    TASK_PIPELINE = 2


class AgentQueueManager:
    """智能体队列管理器"""
    user_id:UUID #对应用户的id
    invoke_from:InvokeFrom #队列
    redis_client:Redis
    _queues:dict[str, Queue]


    def __init__(self, user_id: UUID,invoke_from:InvokeFrom):
        """队列管理器构造函数"""
        #1、初始化数据
        self.user_id = user_id
        self.invoke_from = invoke_from
        self._queues = {}

        #2、内部初始化clinent
        from app.http.module import app_injector
        self.redis_client = app_injector.get(Redis)

    def listen(self, task_id:UUID)->Generator:
        """监听队列函数"""
        #1、定义基础的数据，记录超时时间， 开始时间 ， 最后一次ping 时间
        listen_timeout = 600
        start_time = time.time()
        last_ping_time = 0

        #2、创建循环列表执行死循环 监听队列数据
        while True:
            try:
                #3、从队列中提取数据并检测数据是否存在， 如果存在则直接返回
                item = self.queue(task_id).get(timeout=1)
                if item is None:
                    break
                yield  item
            except queue.Empty:
                continue
            finally:
                #4、计算获取数据的 总耗时， 如果超时则往队列里面 添加停止事件
                elapsed_time = time.time() - start_time
                if elapsed_time >= listen_timeout:
                    pass
                #5、每10s 发起一次ping 请求避免接口中断
                if elapsed_time//10 > last_ping_time:
                    self.publish(task_id,AgentThought(
                        id = uuid.uuid4(),
                        task_id=task_id,
                        event=QueueEvent.PING
                    ))

                    last_ping_time = elapsed_time//10
                #6、判断总耗时是否已经超时
                if elapsed_time >= listen_timeout:
                    self.publish(task_id, AgentThought(
                        id=uuid.uuid4(),
                        task_id=task_id,
                        event=QueueEvent.TIMEOUT
                    ))

                #7、检测是否停止
                if self._is_stopped(task_id):
                    self.publish(task_id, AgentThought(
                        id=uuid.uuid4(),
                        task_id=task_id,
                        event=QueueEvent.STOP
                    ))


    def stop_listen(self, task_id:UUID)->None:
        """停止监听队列"""
        self.queue(task_id).put(None)


    def publish(self, task_id:UUID,agent_queue_event: AgentThought):
        """发布事件"""
        #1、将事件添加到队列中
        self.queue(task_id).put(agent_queue_event)

        #2、检测事件类型是否为需要停止队列的类型， 涵盖 STOP  ERROR TIMEOUT AGENT_END
        if agent_queue_event.event in [QueueEvent.STOP, QueueEvent.ERROR, QueueEvent.TIMEOUT, QueueEvent.AGENT_END]:
            self.stop_listen(task_id)


    def publish_error(self, task_id:UUID, error):
        """发布错误事件"""
        self.publish(task_id, AgentThought(
            id=uuid.uuid4(),
            task_id=task_id,
            event=QueueEvent.ERROR,
            observation=str(error)
        ))


    def _is_stopped(self, task_id:UUID)->bool:
        """检查任务是否停止"""
        task_stopped_cache_key = self.generate_task_stop_cache_key(task_id)
        result = self.redis_client.get(task_stopped_cache_key)
        if result is not None:
            return True
        else:
            return False

    @classmethod
    def set_stop_flag(cls, user_id: UUID, task_id:UUID, invoke_from:InvokeFrom)->None:
        """设置任务停止标志"""
        from app.http.module import app_injector
        redis_client = app_injector.get(Redis)


        #1、检测任务是否存在 如果不存在则 直接返回
        result: bytes = redis_client.get(cls.generate_task_belong_cache_key(task_id))
        if result is None:
            return
        #2、查询当前任务和用户的信息是否匹配， 如果不匹配则直接结束
        user_prefix = "account" if invoke_from in [InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API,InvokeFrom.DEBUGGER] else "end-user"
        if result.decode("utf-8") != f"{user_prefix}-{str(user_id)}":
            return
        #3、构建停止缓存键， 并设置缓存记录
        task_stopped_cache_key = cls.generate_task_stop_cache_key(task_id)
        redis_client.setex(task_stopped_cache_key, 600, "1")



    @classmethod
    def generate_task_belong_cache_key(cls, task_id: UUID)-> str:
        """生成任务的专属缓存键"""
        return f"generate_task_belong:{str(task_id)}"



    @classmethod
    def generate_task_stop_cache_key(cls, task_id:UUID)->str:
        """生成 任务停止键"""
        return f"generate_task_stopped:{str(task_id)}"


    def queue(self, task_id: UUID)->Queue:
        """根据传入的task_id获取对应的任务队列信息"""
        #1、冲队列字典中获取对应的任务队列信息
        q = self._queues.get(str(task_id))
        #2、检测队列是否存在， 如果不存在则闯将队列
        if not q:
            # 3.添加缓存键标识
            user_prefix = "account" if self.invoke_from in [
                InvokeFrom.WEB_APP, InvokeFrom.DEBUGGER, InvokeFrom.ASSISTANT_AGENT,
            ] else "end-user"

            # 4.设置任务对应的缓存键，代表这次任务已经开始了
            self.redis_client.setex(
                self.generate_task_belong_cache_key(task_id),
                1800,
                f"{user_prefix}-{str(self.user_id)}",
            )
            # 5、将队列任务添加到字典中
            q = Queue()
            self._queues[str(task_id)] = q
        return q
